import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import os


def show_bmdl_img(x, y, num, path):
    plt.plot(x, y)
    ax = plt.gca()  # 获取到当前坐标轴信息
    ax.xaxis.set_ticks_position('top')  # 将X坐标轴移到上面
    ax.invert_yaxis()  # 反转Y坐标轴
    plt.savefig(r'{}/{}.png'.format(path, num))
    plt.close()


def show_dat_img(x, y, num, path):
    plt.figure(figsize=(12, 4))
    plt.plot(x, y, 'ro')
    plt.savefig(r'{}/{}.png'.format(path, num))
    plt.close()


# 保存预测频散曲线图像
def show_pre_img(x, y, pre_y, epoch, path):
    """
    :param x: 输入的x
    :param y: 输入的y值
    :param pre_y: 预测的值[batch,c,w]
    :param epoch: 打印的第几个batch 用于命名
    :param path: 保存路径
    :return:
    """
    # 准备工作 检查路径是否存在
    # 保存在传入path的picture目录下
    path = os.path.join(path, 'picture')
    if not os.path.exists(path):
        os.makedirs(path)

    # 修改为其他数据集时 记得更改pre_x的导入路径
    pre_x = pd.read_table(r'D:\model_100\valid\surfwav8501.dat', sep='\t', header=None)[0].to_numpy()

    x = (x.cpu().squeeze().numpy() + 1) * 2300 / 2 + 200
    y = y.cpu().squeeze().numpy() * 1000
    pre_y = pre_y.cpu().squeeze().numpy() * 1000
    for i in range(pre_y.shape[0]):
        #show_bmdl_img(np.arange(x.shape[0]), x[i], str(pre_y.shape[0] * epoch + i) + '_x', path)
        show_dat_img(pre_x, y[i], str(pre_y.shape[0] * epoch + i) + '_y', path)
        show_dat_img(pre_x, pre_y[i], str(pre_y.shape[0] * epoch + i) + '_y_pre', path)


if __name__ == '__main__':
    pass
